import numpy as np
import matplotlib.pyplot as plt
import os
import time
import pandas as pd


class tac_info:
    def __init__(self, grasp_force_id):
        self.tacdata_list = []
        self.grasp_force_list = []
        self.predicted_delta_list = []
        self.time_list = []
        self.grasp_force_id = grasp_force_id  # Tac3D对应的一维力传感器序号


class ExpDataProcessor:
    def __init__(self, SN1, SN2, SN_gf_dict):
        """
        Initializes the ExpDataProcessor with data and an option for online training.

        Args:
            data: The input data to be processed.
            online_training (bool): Flag indicating if the processor is used for online training.
        """

        self.Tac3D_SN1 = SN1
        self.Tac3D_SN2 = SN2
        self.SN_gf_dict = SN_gf_dict  # Dictionary mapping sensor IDs to grasp force IDs
        self.tac_info_dict = {
            self.Tac3D_SN1: tac_info(SN_gf_dict[self.Tac3D_SN1]),
            self.Tac3D_SN2: tac_info(SN_gf_dict[self.Tac3D_SN2]),
        }

        self.dataset_dir = "datalib"
        self.dataset_name = self._get_dataset_dir()  # Get the dataset directory
        if not os.path.exists(self.dataset_name):
            os.makedirs(self.dataset_name)
        print(f"Dataset directory set to: {self.dataset_name}")

    def add_data(self, SN, time_stamp, tacdata, predicted_delta, grasp_force):
        """
        Adds a new set of data to the processor.

        Args:
            tacdata(np.ndarray): The tactile data.
            predicted_delta(np.ndarray): The predicted delta values.
        """

        self.tac_info_dict[SN].tacdata_list.append(tacdata)
        self.tac_info_dict[SN].predicted_delta_list.append(predicted_delta)
        self.tac_info_dict[SN].grasp_force_list.append(grasp_force)
        self.tac_info_dict[SN].time_list.append(time_stamp)

    def show_data(self):
        """
        Displays the current data in the processor.
        """
        for SN, info in self.tac_info_dict.items():
            force_data = np.sum(np.stack(info.tacdata_list)[..., 6:9], axis=1)  # Sum the forces across all contacts
            tangential_force = np.linalg.norm(force_data[:, 0:2], axis=1)
            normal_force = -force_data[:, 2]
            predicted_delta = np.array(info.predicted_delta_list)
            time_list = np.array(info.time_list) - info.time_list[0]
            grasp_force = np.array(info.grasp_force_list)

            plt.figure(figsize=(12, 8))
            plt.subplot(2, 1, 1)
            plt.plot(time_list, tangential_force, label="Tangential Force", color="blue")
            plt.plot(time_list, normal_force, label="Normal Force", color="red")
            plt.plot(time_list, grasp_force, label="Grasp Force", color="green")
            plt.title(f"Tactile Data for {SN}")
            plt.ylabel("Force (N)")
            plt.legend()

            plt.subplot(2, 1, 2)
            plt.plot(time_list, predicted_delta, label="Predicted Delta", color="purple")
            plt.xlabel("Time (s)")
            plt.ylabel("Predicted Delta")
            plt.legend()
            plt.tight_layout()
            plt.show()

    def get_mae(self, goal_delta: float, exp_result: bool):
        mae_dict = {}
        est_mae = 0.0
        track_mae = 0.0
        for SN, info in self.tac_info_dict.items():
            tacdata = np.stack(info.tacdata_list)
            predicted_delta = np.array(info.predicted_delta_list)
            true_delta = self._get_true_delta(exp_result, tacdata)
            mae = np.mean(np.abs(predicted_delta - true_delta))
            mae_dict[f"{SN}_est_mae"] = mae
            est_mae += mae
            mae = np.mean(np.abs(predicted_delta - goal_delta))
            mae_dict[f"{SN}_track_mae"] = mae
            track_mae += mae
        return mae_dict, est_mae / len(self.tac_info_dict), track_mae / len(self.tac_info_dict)

    def save_data(self, exp_result: bool):
        dir = f"{self.dataset_name}/{time.strftime('%Y%m%d_%H%M%S')}"
        if not os.path.exists(dir):
            os.makedirs(dir)
        for SN, info in self.tac_info_dict.items():
            tacdata = np.stack(info.tacdata_list)
            grasp_force = np.array(info.grasp_force_list)
            predicted_delta = np.array(info.predicted_delta_list)
            true_delta = self._get_true_delta(exp_result, tacdata)

            tacdata_file = os.path.join(dir, f"tacdata_{self.SN_gf_dict[SN]}.npy")
            np.save(tacdata_file, tacdata)
            grasp_force_file = os.path.join(dir, f"grasp_force_{self.SN_gf_dict[SN]}.npy")
            np.save(grasp_force_file, grasp_force)
            predicted_delta_file = os.path.join(dir, f"predicted_delta_{self.SN_gf_dict[SN]}.npy")
            np.save(predicted_delta_file, predicted_delta)
            true_delta_file = os.path.join(dir, f"true_delta_{self.SN_gf_dict[SN]}.npy")
            np.save(true_delta_file, true_delta)

        print(f"Data saved to {dir}")
        return dir

    def add_actor_data(self, dir, actor_dict: dict):
        """
        Adds actor data to the processor.
        This method is a placeholder for future implementation.
        """
        if len(actor_dict) == 0:
            return
        data = pd.DataFrame(actor_dict)
        data.to_csv(os.path.join(dir, "actor_data.csv"), index=True)
        print(f"Saved actor data to {dir}")

    def clear_data(self):
        """
        Clears the current data in the processor.
        """
        for SN in self.tac_info_dict:
            self.tac_info_dict[SN].tacdata_list.clear()
            self.tac_info_dict[SN].grasp_force_list.clear()
            self.tac_info_dict[SN].predicted_delta_list.clear()
            self.tac_info_dict[SN].time_list.clear()

        print("Data cleared.")

    def _get_true_delta(self, exp_result, tacdata):
        """
        Calculates the true delta based on the experimental result and tactile data.
        Args:
            exp_result (bool): The result of the experiment.
            tacdata (np.ndarray): The tactile data.
        """
        sensored_force = np.sum(tacdata[:, :, 6:9], axis=1)
        tangent_force = np.linalg.norm(sensored_force[:, 0:2], axis=1)
        normal_force = -sensored_force[:, 2]

        true_delta = np.zeros_like(normal_force)
        if exp_result == True:  # exp success, the last value means slip
            goal_idx = -1
            while abs(normal_force[goal_idx] - normal_force[goal_idx - 1]) > 1:
                goal_idx -= 1
            final_normal_force = normal_force[goal_idx]
            final_tangent_force = tangent_force[goal_idx]
        else:  # exp failed, the maximum tangent force is the slip point
            final_normal_force = normal_force[np.argmax(tangent_force)]
            final_tangent_force = np.max(tangent_force)

        true_delta = 1 - (tangent_force / (final_tangent_force + 1e-8)) / (
            normal_force / (final_normal_force + 1e-8) + 1e-4
        )

        true_delta = np.clip(true_delta, 0, 1)  # Ensure delta is between 0 and 1
        return true_delta

    def _get_dataset_dir(self):
        """
        Returns the directory where the dataset is stored.
        Data should be stored in a directory named 'dataset_{index}' where index is an integer.
        """
        dataset_idx = 1
        dataset_path = f"{self.dataset_dir}/dataset_{dataset_idx}"
        while os.path.exists(dataset_path):
            used_file_flag = os.path.join(dataset_path, "used_dataset")
            if os.path.exists(used_file_flag):
                dataset_idx += 1
                dataset_path = f"{self.dataset_dir}/dataset_{dataset_idx}"
            else:
                break
        return dataset_path
